#include <queue>
#include <vector>
#include <arm_neon.h>
#include <algorithm>
#include <random>
#include <numeric>
#include <mutex>
#include <omp.h>

struct IVF_PQ_Index {
    std::vector<std::vector<float>> centroids;  // IVF聚类中心
    std::vector<std::vector<uint32_t>> lists;   // 倒排列表
    size_t nlist = 32;                         // 聚类中心数量
    size_t sub_dim;                             // PQ子空间维度
    size_t m = 4;                               // PQ子空间数
    size_t ks = 64;                            // PQ每子空间簇数
    std::vector<std::vector<float>> pq_codebooks;
    std::vector<std::vector<uint8_t>> pq_codes;
    bool initialized = false;
};

static IVF_PQ_Index ivf_pq_index;
static std::mutex init_mutex;

float compute_distance_neon(const float* x, const float* y, size_t dim) {
    float32x4_t sum = vdupq_n_f32(0.0f);
    size_t d;
    for (d = 0; d + 4 <= dim; d += 4) {
        float32x4_t x_vec = vld1q_f32(x + d);
        float32x4_t y_vec = vld1q_f32(y + d);
        float32x4_t diff = vsubq_f32(x_vec, y_vec);
        sum = vmlaq_f32(sum, diff, diff);
    }
    float32x2_t sum2 = vadd_f32(vget_low_f32(sum), vget_high_f32(sum));
    float result = vget_lane_f32(sum2, 0) + vget_lane_f32(sum2, 1);
    
    // 处理剩余部分
    for (; d < dim; ++d) {
        float diff = x[d] - y[d];
        result += diff * diff;
    }
    return result;
}

void train_pq_codebook(const float* base, size_t base_number, size_t vecdim) {
    ivf_pq_index.sub_dim = vecdim / ivf_pq_index.m;
    ivf_pq_index.pq_codebooks.resize(ivf_pq_index.m);
    ivf_pq_index.pq_codes.resize(base_number, std::vector<uint8_t>(ivf_pq_index.m));

    #pragma omp parallel for
    for (size_t k = 0; k < ivf_pq_index.m; ++k) {
        // 分配空间
        ivf_pq_index.pq_codebooks[k].resize(ivf_pq_index.ks * ivf_pq_index.sub_dim);
        
        // 提取子向量
        std::vector<std::vector<float>> sub_vectors(base_number);
        for (size_t i = 0; i < base_number; ++i) {
            sub_vectors[i].resize(ivf_pq_index.sub_dim);
            std::copy_n(base + i * vecdim + k * ivf_pq_index.sub_dim, 
                       ivf_pq_index.sub_dim, 
                       sub_vectors[i].data());
        }

        // K-means聚类生成codebook
        std::vector<std::vector<float>> cluster_centers(ivf_pq_index.ks);
        std::vector<size_t> assignments(base_number);
        
        // 随机初始化聚类中心
        std::random_device rd;
        std::mt19937 gen(rd());
        std::shuffle(sub_vectors.begin(), sub_vectors.end(), gen);
        for (size_t j = 0; j < ivf_pq_index.ks; ++j) {
            cluster_centers[j] = sub_vectors[j];
        }

        // K-means迭代
        for (size_t iter = 0; iter < 10; ++iter) {
            // 分配点到最近的聚类
            #pragma omp parallel for
            for (size_t i = 0; i < base_number; ++i) {
                float min_dist = std::numeric_limits<float>::max();
                size_t best_cluster = 0;
                for (size_t j = 0; j < ivf_pq_index.ks; ++j) {
                    float dist = compute_distance_neon(sub_vectors[i].data(), 
                                                     cluster_centers[j].data(), 
                                                     ivf_pq_index.sub_dim);
                    if (dist < min_dist) {
                        min_dist = dist;
                        best_cluster = j;
                    }
                }
                assignments[i] = best_cluster;
            }

            // 重新计算聚类中心
            std::vector<size_t> counts(ivf_pq_index.ks, 0);
            std::vector<std::vector<float>> sums(ivf_pq_index.ks, 
                                               std::vector<float>(ivf_pq_index.sub_dim, 0.0f));
            
            #pragma omp parallel for
            for (size_t i = 0; i < base_number; ++i) {
                size_t cluster = assignments[i];
                #pragma omp atomic
                counts[cluster]++;
                for (size_t d = 0; d < ivf_pq_index.sub_dim; ++d) {
                    #pragma omp atomic
                    sums[cluster][d] += sub_vectors[i][d];
                }
            }
            
            #pragma omp parallel for
            for (size_t j = 0; j < ivf_pq_index.ks; ++j) {
                if (counts[j] > 0) {
                    for (size_t d = 0; d < ivf_pq_index.sub_dim; ++d) {
                        cluster_centers[j][d] = sums[j][d] / counts[j];
                    }
                }
            }
        }

        // 存储最终的codebook
        #pragma omp parallel for
        for (size_t j = 0; j < ivf_pq_index.ks; ++j) {
            std::copy_n(cluster_centers[j].data(), 
                       ivf_pq_index.sub_dim, 
                       ivf_pq_index.pq_codebooks[k].data() + j * ivf_pq_index.sub_dim);
        }

        // 使用学习到的codebook编码所有向量
        #pragma omp parallel for
        for (size_t i = 0; i < base_number; ++i) {
            float min_dist = std::numeric_limits<float>::max();
            uint8_t best_codeword = 0;
            for (uint8_t j = 0; j < ivf_pq_index.ks; ++j) {
                const float* codebook_ptr = ivf_pq_index.pq_codebooks[k].data() + j * ivf_pq_index.sub_dim;
                float dist = compute_distance_neon(sub_vectors[i].data(), codebook_ptr, ivf_pq_index.sub_dim);
                if (dist < min_dist) {
                    min_dist = dist;
                    best_codeword = j;
                }
            }
            ivf_pq_index.pq_codes[i][k] = best_codeword;
        }
    }
}

void build_ivf_index(const float* base, size_t base_number, size_t vecdim) {
    ivf_pq_index.centroids.resize(ivf_pq_index.nlist);
    
    // 1. 随机选择初始中心
    std::vector<size_t> indices(base_number);
    std::iota(indices.begin(), indices.end(), 0);
    std::random_device rd;
    std::mt19937 gen(rd());
    std::shuffle(indices.begin(), indices.end(), gen);
    
    #pragma omp parallel for
    for (size_t i = 0; i < ivf_pq_index.nlist; ++i) {
        ivf_pq_index.centroids[i].resize(vecdim);
        std::copy_n(base + indices[i] * vecdim, vecdim, ivf_pq_index.centroids[i].data());
    }

    // 2. 构建倒排列表
    ivf_pq_index.lists.resize(ivf_pq_index.nlist);
    #pragma omp parallel for
    for (size_t i = 0; i < base_number; ++i) {
        float min_dist = std::numeric_limits<float>::max();
        size_t best_cluster = 0;
        const float* vec = base + i * vecdim;
        
        for (size_t j = 0; j < ivf_pq_index.nlist; ++j) {
            float dist = compute_distance_neon(vec, ivf_pq_index.centroids[j].data(), vecdim);
            if (dist < min_dist) {
                min_dist = dist;
                best_cluster = j;
            }
        }
        
        #pragma omp critical
        ivf_pq_index.lists[best_cluster].push_back(i);
    }
}

void init_ivf_pq(const float* base, size_t base_number, size_t vecdim) {
    std::lock_guard<std::mutex> lock(init_mutex);
    if (ivf_pq_index.initialized) return;

    // 先训练PQ编码
    train_pq_codebook(base, base_number, vecdim);
    
    // 然后构建IVF索引
    build_ivf_index(base, base_number, vecdim);

    ivf_pq_index.initialized = true;
}

std::priority_queue<std::pair<float, uint32_t>, 
                   std::vector<std::pair<float, uint32_t>>, 
                   std::greater<std::pair<float, uint32_t>>> 
openmp_pq_search(const float* base, const float* query, size_t base_number, 
             size_t vecdim, size_t top_k, size_t nprobe = 10) 
{
    if (!ivf_pq_index.initialized) {
        init_ivf_pq(base, base_number, vecdim);
    }

    nprobe = std::min(nprobe, ivf_pq_index.nlist);
    
    // 1. 粗搜索（找出最近的nprobe个聚类中心）
    std::vector<std::pair<float, size_t>> centroid_dists(ivf_pq_index.nlist);
    #pragma omp parallel for
    for (size_t i = 0; i < ivf_pq_index.nlist; ++i) {
        float dist = compute_distance_neon(query, ivf_pq_index.centroids[i].data(), vecdim);
        centroid_dists[i] = {dist, i};
    }
    
    // 部分排序找出前nprobe个最小的
    std::partial_sort(centroid_dists.begin(), centroid_dists.begin() + nprobe, 
                     centroid_dists.end());

    // 2. 收集候选向量
    std::vector<uint32_t> candidates;
    for (size_t i = 0; i < nprobe; ++i) {
        size_t cluster_id = centroid_dists[i].second;
        const auto& list = ivf_pq_index.lists[cluster_id];
        candidates.insert(candidates.end(), list.begin(), list.end());
    }

    // 3. 精确计算距离
    std::vector<float> dists(candidates.size());
    #pragma omp parallel for
    for (size_t i = 0; i < candidates.size(); ++i) {
        uint32_t idx = candidates[i];
        const float* vec = base + idx * vecdim;
        dists[i] = compute_distance_neon(query, vec, vecdim);
    }

    // 4. 构建优先队列
    std::priority_queue<std::pair<float, uint32_t>, 
                       std::vector<std::pair<float, uint32_t>>, 
                       std::greater<std::pair<float, uint32_t>>> q;
    
    for (size_t i = 0; i < candidates.size(); ++i) {
        if (q.size() < top_k) {
            q.emplace(dists[i], candidates[i]);
        } else if (dists[i] < q.top().first) {
            q.pop();
            q.emplace(dists[i], candidates[i]);
        }
    }
    
    return q;
}